package com.example.demo.ws;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import java.util.Map;

@Component
public class WsInterceptor extends HttpSessionHandshakeInterceptor {
    private static final Logger log = LoggerFactory.getLogger(WsInterceptor.class);

    static final String ATTR_KEY_USER_ID = "userId";

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        log.info("beforeHandshake uri: {}", request.getURI());
        String auth = ((ServletServerHttpRequest) request).getServletRequest().getParameter("auth");
        log.info("auth: {}", auth);
        if (auth == null || auth.trim().length() == 0) {
            log.info("Handshake failure, auth is empty", auth);
            return false;
        }

        Long userId = getUserIdByAuth(auth);
        if (userId == null) {
            log.info("Handshake failure, auth: {} error", auth);
        } else {
            attributes.put(ATTR_KEY_USER_ID, userId);
        }
        return userId != null;
    }

    private static Long getUserIdByAuth(String auth) {
        if (auth == null) {
            return null;
        }
        switch (auth) {
            case "aa":
                return 1L;
            case "bb":
                return 2L;
            case "cc":
                return 3L;
            case "dd":
                return 4L;
            case "ee":
                return 5L;
            default:
                return null;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception ex) {
        log.info("afterHandshake");
        super.afterHandshake(request, response, wsHandler, ex);
    }
}